use std::{collections::HashMap, io::Write, path::Path, process::Command};

use anyhow::{Context, Result, bail, ensure};
use cargo_semver_checks::ReleaseType;
use clap::Args;
use esp_metadata::Chip;
use serde::{Deserialize, Serialize};
use strum::IntoEnumIterator;
use toml_edit::{Item, Value};

use crate::{
    Package,
    cargo::CargoToml,
    commands::{VersionBump, checker::min_package_update, do_version_bump},
    git::current_branch,
};

/// Arguments for generating a release plan.
#[derive(Debug, Args)]
pub struct PlanArgs {
    /// Allow making a release from the current (non-main) branch. The
    /// branch is expected to exist in the upstream esp-hal repo.
    #[arg(long)]
    allow_non_main: bool,

    /// The packages to be released.
    #[arg(value_enum, default_values_t = Package::iter())]
    packages: Vec<Package>,
}

/// A package in the release plan.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PackagePlan {
    pub package: Package,
    pub semver_checked: bool,
    pub current_version: semver::Version,
    pub new_version: semver::Version,
    pub tag_name: String,
    /// The version bump that will be applied to the package.
    pub bump: VersionBump,
}

/// A release plan is a list of packages and their version increments.
///
/// The plan is generated by analyzing the workspace and its dependencies.
/// The order of the packages in the plan is important, as it determines the
/// order in which the packages are released.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Plan {
    /// The branch the release is made from.
    pub base: String,
    /// The packages to be released, in order.
    pub packages: Vec<PackagePlan>,
}

impl Plan {
    /// Load a release plan from a file.
    pub fn from_path(plan_path: &Path) -> Result<Self> {
        let plan_source = std::fs::read_to_string(&plan_path)
            .with_context(|| format!("Failed to read release plan from {}. Run `cargo xrelease plan` to generate a release plan.", plan_path.display()))?;

        if plan_source.lines().any(|line| line.starts_with("//")) {
            bail!(
                "The release plan has not been finalized. Please open the plan and follow the instructions in it."
            );
        }

        let plan = serde_json::from_str::<Plan>(&plan_source).with_context(|| {
            format!("Failed to parse release plan from {}", plan_path.display())
        })?;

        Ok(plan)
    }
}

/// Generate a release plan for the specified packages.
pub fn plan(workspace: &Path, args: PlanArgs) -> Result<()> {
    let current_branch = ensure_main_branch(args.allow_non_main)?;

    // Recursively collect dependencies. A bit inefficient, but we don't need to
    // sort a lot.
    let mut packages_to_release = args
        .packages
        .iter()
        .filter(|p| p.is_published())
        .flat_map(|p| related_crates(workspace, *p))
        .collect::<Vec<_>>();

    packages_to_release.sort();
    packages_to_release.dedup();

    let mut package_tomls = packages_to_release
        .iter()
        .map(|pkg| CargoToml::new(workspace, *pkg).map(|cargo_toml| (*pkg, cargo_toml)))
        .collect::<Result<HashMap<_, _>>>()?;

    // Determine package dependencies (package -> dependencies)
    let mut dep_graph = HashMap::new();
    for (package, toml) in package_tomls.iter_mut() {
        dep_graph.insert(*package, toml.repo_dependencies());
    }

    // Topological sort the packages into a release order. Note that this is not a stable order,
    // because the source data is a HashMap which does not have a stable insertion order. This is
    // okay, as long as the relationships of the dependencies are kept.
    let sorted = topological_sort(&dep_graph);
    log::debug!("Sorted packages: {:?}", sorted);

    for package in sorted.iter() {
        ensure!(
            package_tomls[&package].is_published(),
            "Cannot release {package}: package is not published. The command \
            would have released the following packages: {}",
            sorted
                .iter()
                .map(|p| p.to_string())
                .collect::<Vec<_>>()
                .join(", ")
        );
    }

    // Gather semver bump requirements
    let mut update_amounts = vec![];
    let all_chips = Chip::iter().collect::<Vec<_>>();

    let changed = collect_changed_packages(
        &sorted,
        // Read direct dependencies out of Cargo.toml
        &|p| CargoToml::new(workspace, p).unwrap().repo_dependencies(),
        &|p| package_changed_since_last_release(workspace, p),
    );

    for package in sorted.iter().copied() {
        let amount = if changed[&package] {
            let amount = if package.is_semver_checked() {
                min_package_update(workspace, package, &all_chips)?
            } else {
                let forever_unstable = if let Some(metadata) =
                    package_tomls[&package].espressif_metadata()
                    && let Some(Item::Value(forever_unstable)) = metadata.get("forever-unstable")
                {
                    // Special case: some packages are perma-unstable, meaning they won't ever have
                    // a stable release. For these packages, we always use a
                    // patch release.
                    if let Value::Boolean(forever_unstable) = forever_unstable {
                        *forever_unstable.value()
                    } else {
                        log::warn!(
                            "Invalid value for 'forever-unstable' in metadata - must be a boolean"
                        );
                        true
                    }
                } else {
                    false
                };

                if forever_unstable {
                    ReleaseType::Patch
                } else {
                    ReleaseType::Minor
                }
            };

            log::debug!("{} needs {:?} version bump", package, amount);
            Some(amount)
        } else {
            log::debug!("{} has no changes since last release", package);
            None
        };

        update_amounts.push((package, amount));
    }

    // Generate plan file. The plan should include, as an ordered list, the packages
    // and their version increment.

    let plan = Plan {
        base: current_branch,
        packages: update_amounts
            .into_iter()
            .filter_map(|(package, bump)| {
                bump.map(|b| {
                    let current_version = package_tomls[&package].package_version();

                    let bump = if !current_version.pre.is_empty() {
                        VersionBump::PreRelease(
                            current_version
                                .pre
                                .as_str()
                                .split('.')
                                .next()
                                .unwrap()
                                .to_string(),
                        )
                    } else {
                        match b {
                            ReleaseType::Major => VersionBump::Major,
                            ReleaseType::Minor => VersionBump::Minor,
                            ReleaseType::Patch => VersionBump::Patch,
                            _ => unreachable!(),
                        }
                    };

                    let new_version = do_version_bump(&current_version, &bump).unwrap();
                    let tag_name = package.tag(&new_version);

                    PackagePlan {
                        package,
                        semver_checked: package.is_semver_checked(),
                        current_version,
                        new_version,
                        tag_name,
                        bump,
                    }
                })
            })
            .collect(),
    };

    let plan_path = workspace.join("release_plan.jsonc");
    let plan_path = crate::windows_safe_path(&plan_path);
    log::debug!("Writing release plan to {}", plan_path.display());

    let plan_header = String::from(
        r#"// Release plan generated by xtask.
//
// This file is automatically generated and meant to be manually edited. Change the version bump
// where appropriate. Remove packages that should not be released. Delete the comment block to
// finalize the release plan, otherwise trying to apply the release plan will fail.
//
// This file contains each package that needs to be released. The order of the packages is
// important, the packages will be released in the order they appear in the release plan.
// This is done in case the tool gets the dependency graph wrong.
//
// This plan may also contain packages that do not need to be released. Please review them
// in case the release was triggered by commits that only included non-relevant changes. These
// packages can be removed from the plan.
//
// You don't need to change versions and tag names. The plan will be updated when applying the
// release plan. The release plan will then be added to the release PR and will be used to
// orchestrate the actual publishing process.
//
// For each package, this plan contains the version bump that will be applied.
// The version bump is one of:
// - PreRelease (`"bump": { "PreRelease": "beta" }`)
// - Patch (`"bump": "Patch"`)
// - Minor (`"bump": "Minor"`)
// - Major (`"bump": "Major"`)
"#,
    );

    let mut plan_file = std::fs::File::create(&plan_path)
        .with_context(|| format!("Failed to create {plan_path:?}"))?;
    plan_file
        .write_all(plan_header.as_bytes())
        .with_context(|| format!("Failed to write to {plan_file:?}"))?;
    serde_json::to_writer_pretty(&mut plan_file, &plan)
        .with_context(|| format!("Failed to serialize {plan:?} as pretty-printed JSON"))?;
    log::debug!("Release plan written to {}", plan_path.display());

    println!("Release plan written to {}.", plan_path.display());
    println!("Please review the release plan and make changes where appropriate.");
    println!(
        "To apply the release plan, you'll need to remove the heading comment, save the \
        file, then run the following command: `cargo xrelease execute-plan`",
    );

    Ok(())
}

fn collect_changed_packages(
    sorted: &[Package],
    direct_dependencies: &impl Fn(Package) -> Vec<Package>,
    is_changed: &impl Fn(Package) -> bool,
) -> HashMap<Package, bool> {
    let mut changed = HashMap::new();

    // Collect recursive dependencies.
    let mut all_dependencies_of = HashMap::new();
    for package in sorted.iter().copied() {
        let dependencies = related_crates_cb(package, direct_dependencies);
        all_dependencies_of.insert(package, dependencies);
    }

    // Check each package whether they should be considered for release. A package should be
    // released if it has changes since the last release, or if it depends on a package that has
    // changes.
    for package in sorted.iter().copied() {
        if changed.get(&package).copied() == Some(true) {
            // This package has already been marked as changed, so skip it.
            continue;
        }

        let is_changed = is_changed(package);
        changed.insert(package, is_changed);
        if is_changed {
            // Mark all direct and indirect _dependents_ as changed too.
            for (dependent, deps) in all_dependencies_of.iter() {
                if deps.contains(&package) {
                    changed.insert(*dependent, true);
                }
            }
        }
    }

    changed
}

/// Ensure we are on the main branch, or allow non-main if specified.
pub fn ensure_main_branch(allow_non_main: bool) -> Result<String> {
    let current_branch = current_branch()?;

    if allow_non_main {
        println!("Release will be made from branch: {current_branch}");
    } else {
        ensure!(
            current_branch == "main",
            "You are not on the main branch. Please switch to the main branch before running this command again, or add `--allow-non-main`."
        );
    }

    Ok(current_branch)
}

fn package_changed_since_last_release(workspace: &Path, package: Package) -> bool {
    let toml = CargoToml::new(workspace, package).expect("Failed to load Cargo.toml");

    let last_version: semver::Version = toml.package_version();

    let last_tag = package.tag(&last_version);
    let mut commits = commits_since_tag(workspace, &toml.package_path(), &last_tag);

    if commits == 0 {
        // Try to look up the last tag in the git history and retry.
        let last_tag = Command::new("git")
            .args(["describe", "--tags", "--abbrev=0"])
            .current_dir(workspace)
            .output()
            .expect("Failed to get last tag");
        let last_tag = String::from_utf8_lossy(&last_tag.stdout);
        let last_tag = last_tag.trim();

        if !last_tag.is_empty() {
            commits = commits_since_tag(workspace, &toml.package_path(), last_tag);
        }
    }

    log::trace!("{package}: {commits} commits since last release");

    commits > 0
}

fn commits_since_tag(workspace: &Path, package_path: &Path, tag: &str) -> usize {
    let commits = Command::new("git")
        .args([
            "log",
            &format!("{tag}..HEAD"),
            "--oneline",
            &package_path.display().to_string(),
        ])
        .current_dir(workspace)
        .output()
        .expect("Failed to get commit count since last tag");

    let output = String::from_utf8_lossy(&commits.stdout);

    output.lines().filter(|l| !l.trim().is_empty()).count()
}

/// Collects dependencies recursively, based on a callback that provides direct dependencies.
///
/// This is an implementation detail of the `related_crates` function, separate for testing
/// purposes.
fn related_crates_cb(
    package: Package,
    direct_dependencies: &impl Fn(Package) -> Vec<Package>,
) -> Vec<Package> {
    let mut packages = vec![package];

    for dep in direct_dependencies(package) {
        for dep in related_crates_cb(dep, direct_dependencies) {
            if !packages.contains(&dep) {
                packages.push(dep);
            }
        }
    }

    packages
}

fn related_crates(workspace: &Path, package: Package) -> Vec<Package> {
    related_crates_cb(package, &|p| {
        CargoToml::new(workspace, p).unwrap().repo_dependencies()
    })
}

fn topological_sort(dep_graph: &HashMap<Package, Vec<Package>>) -> Vec<Package> {
    let mut sorted = Vec::new();
    let mut dep_graph = dep_graph.clone();
    while !dep_graph.is_empty() {
        dep_graph.retain(|pkg, deps| {
            deps.retain(|dep| !sorted.contains(dep));

            if deps.is_empty() {
                sorted.push(*pkg);
                false
            } else {
                true
            }
        });
    }

    sorted
}

#[cfg(test)]
mod tests {
    use super::*;

    // Dependencies in this graph do not always reflect real dependencies, it is only a test of
    // correct function operation.
    fn sample_dep_graph() -> HashMap<Package, Vec<Package>> {
        //           |-------------+-> esp-hal -> esp-alloc
        // esp-rtos -+> esp-radio -|

        let mut dep_graph = HashMap::new();
        dep_graph.insert(Package::EspHal, vec![Package::EspAlloc]);
        dep_graph.insert(Package::EspRtos, vec![Package::EspHal, Package::EspRadio]);
        dep_graph.insert(Package::EspRadio, vec![Package::EspHal]);
        dep_graph.insert(Package::EspAlloc, vec![]);

        dep_graph
    }

    #[track_caller]
    fn assert_unstable_lists_equal(actual: &[Package], expected: &[Package]) {
        use std::collections::HashSet;
        let actual_set: HashSet<_> = actual.iter().collect();
        let expected_set: HashSet<_> = expected.iter().collect();
        assert_eq!(actual_set, expected_set);
    }

    #[test]
    fn test_topological_sort() {
        let dep_graph = sample_dep_graph();
        let sorted = topological_sort(&dep_graph);

        assert_eq!(
            sorted,
            [
                Package::EspAlloc,
                Package::EspHal,
                Package::EspRadio,
                Package::EspRtos
            ]
        );
    }

    #[test]
    fn test_change_detection() {
        let dep_graph = sample_dep_graph();
        let sorted = topological_sort(&dep_graph);

        let assert_changes = |changed: &[Package], expected: &[Package]| {
            let changed = collect_changed_packages(&sorted, &|p| dep_graph[&p].clone(), &|p| {
                changed.contains(&p)
            });
            let changed = changed
                .into_iter()
                .filter_map(|(p, ch)| if ch { Some(p) } else { None })
                .collect::<Vec<_>>();

            assert_unstable_lists_equal(&changed, expected);
        };

        assert_changes(&[], &[]);
        assert_changes(&[Package::EspRtos], &[Package::EspRtos]);
        assert_changes(&[Package::EspRadio], &[Package::EspRadio, Package::EspRtos]);
        assert_changes(
            &[Package::EspHal],
            &[Package::EspHal, Package::EspRadio, Package::EspRtos],
        );
    }
}
